package com.ds.lens.flink.monitor.forecast.function;

import com.google.gson.Gson;
import com.ds.lens.data.common.constant.Constants;
import com.ds.lens.data.common.document.ForecastData;
import com.ds.lens.data.common.document.ForecastKafkaData;
import com.ds.lens.flink.monitor.MonitorJobContext;
import com.ds.lens.flink.monitor.forecast.ForecastConfig;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.math.NumberUtils;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch.RequestIndexer;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.Requests;
import org.elasticsearch.common.xcontent.XContentType;

import java.nio.charset.StandardCharsets;

/**
 * Description:
 *
 * @author WeiShaoying
 * @date 2020-05-10
 */
@Slf4j
public class ForecastDataSinkFunction implements ElasticsearchSinkFunction<ForecastKafkaData> {

    private static final long serialVersionUID = -3504570282998478486L;

    private IndexRequest createIndexRequest(ForecastData element) {
        try {
            ForecastConfig forecastConfig = MonitorJobContext.getInstance().getForecastConfig();
            Gson gson = MonitorJobContext.getInstance().getGson();
            return Requests.indexRequest()
                    .index(forecastConfig.getEsForecastDataIndex())
                    .type(forecastConfig.getEsForecastDataType())
                    .id(element.getSubTopic() + element.getTs())
                    .source(gson.toJson(element).getBytes(StandardCharsets.UTF_8), XContentType.JSON);
        } catch (Exception e) {
            log.error("ForecastDataSinkFunction sink failed.", e);
            return null;
        }
    }

    @Override
    public void process(ForecastKafkaData forecastKafkaData, RuntimeContext runtimeContext, RequestIndexer requestIndexer) {
        String subTopic = forecastKafkaData.getSubTopic();
        forecastKafkaData.getData().forEach(forecastData -> {
            forecastData.setSubTopic(subTopic);
            String[] strings = subTopic.split(Constants.DASH);
            if (strings.length >= 2 && NumberUtils.isNumber(strings[0]) && NumberUtils.isNumber(strings[1])) {
                forecastData.setMonitorId(Long.parseLong(strings[0]));
                forecastData.setVersion(Integer.parseInt(strings[1]));

                IndexRequest request = createIndexRequest(forecastData);
                if (request != null) {
                    requestIndexer.add(request);
                }
            }
        });
    }
}
